"""
可视化工具模块

提供数据可视化、结果展示、图表生成等功能
"""

import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
import cv2
from typing import List, Optional, Tuple, Union, Dict, Any
from pathlib import Path
import seaborn as sns
from PIL import Image, ImageDraw, ImageFont
import torch


def setup_matplotlib():
    """设置matplotlib中文字体"""
    plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False
    plt.style.use('seaborn-v0_8')


def plot_image_with_mask(image: np.ndarray, 
                        mask: np.ndarray,
                        title: str = "Image with Mask",
                        alpha: float = 0.5,
                        save_path: Optional[Union[str, Path]] = None) -> None:
    """
    绘制带掩码的图像
    
    Args:
        image: 输入图像 (H, W, C) 或 (H, W)
        mask: 掩码 (H, W)
        title: 图像标题
        alpha: 掩码透明度
        save_path: 保存路径
    """
    setup_matplotlib()
    
    fig, ax = plt.subplots(1, 1, figsize=(10, 8))
    
    # 显示原图
    if len(image.shape) == 3:
        ax.imshow(image)
    else:
        ax.imshow(image, cmap='gray')
    
    # 叠加掩码
    ax.imshow(mask, alpha=alpha, cmap='Reds')
    
    ax.set_title(title, fontsize=14)
    ax.axis('off')
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    plt.show()


def plot_training_curves(train_losses: List[float],
                       val_losses: Optional[List[float]] = None,
                       train_metrics: Optional[Dict[str, List[float]]] = None,
                       val_metrics: Optional[Dict[str, List[float]]] = None,
                       save_path: Optional[Union[str, Path]] = None) -> None:
    """
    绘制训练曲线
    
    Args:
        train_losses: 训练损失
        val_losses: 验证损失
        train_metrics: 训练指标
        val_metrics: 验证指标
        save_path: 保存路径
    """
    setup_matplotlib()
    
    n_plots = 1
    if val_losses is not None:
        n_plots += 1
    if train_metrics is not None or val_metrics is not None:
        n_plots += 1
    
    fig, axes = plt.subplots(1, n_plots, figsize=(5 * n_plots, 4))
    if n_plots == 1:
        axes = [axes]
    
    # 损失曲线
    ax = axes[0]
    epochs = range(1, len(train_losses) + 1)
    ax.plot(epochs, train_losses, 'b-', label='训练损失', linewidth=2)
    
    if val_losses is not None:
        ax.plot(epochs, val_losses, 'r-', label='验证损失', linewidth=2)
    
    ax.set_xlabel('Epoch')
    ax.set_ylabel('Loss')
    ax.set_title('训练曲线')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    # 指标曲线
    if n_plots > 1 and (train_metrics is not None or val_metrics is not None):
        ax = axes[1]
        
        if train_metrics is not None:
            for metric_name, values in train_metrics.items():
                ax.plot(epochs, values, label=f'训练 {metric_name}', linewidth=2)
        
        if val_metrics is not None:
            for metric_name, values in val_metrics.items():
                ax.plot(epochs, values, label=f'验证 {metric_name}', linewidth=2)
        
        ax.set_xlabel('Epoch')
        ax.set_ylabel('Metric Value')
        ax.set_title('指标曲线')
        ax.legend()
        ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    plt.show()


def plot_comparison_grid(images: List[np.ndarray],
                        titles: List[str],
                        masks: Optional[List[np.ndarray]] = None,
                        predictions: Optional[List[np.ndarray]] = None,
                        save_path: Optional[Union[str, Path]] = None) -> None:
    """
    绘制对比网格图
    
    Args:
        images: 图像列表
        titles: 标题列表
        masks: 真实掩码列表
        predictions: 预测掩码列表
        save_path: 保存路径
    """
    setup_matplotlib()
    
    n_images = len(images)
    n_cols = min(4, n_images)
    n_rows = (n_images + n_cols - 1) // n_cols
    
    fig, axes = plt.subplots(n_rows, n_cols, figsize=(4 * n_cols, 4 * n_rows))
    if n_rows == 1:
        axes = axes.reshape(1, -1)
    elif n_cols == 1:
        axes = axes.reshape(-1, 1)
    
    for i in range(n_images):
        row = i // n_cols
        col = i % n_cols
        ax = axes[row, col]
        
        # 显示图像
        if len(images[i].shape) == 3:
            ax.imshow(images[i])
        else:
            ax.imshow(images[i], cmap='gray')
        
        # 叠加掩码
        if masks is not None and i < len(masks):
            ax.imshow(masks[i], alpha=0.5, cmap='Reds')
        
        # 叠加预测
        if predictions is not None and i < len(predictions):
            ax.imshow(predictions[i], alpha=0.3, cmap='Blues')
        
        ax.set_title(titles[i], fontsize=12)
        ax.axis('off')
    
    # 隐藏多余的子图
    for i in range(n_images, n_rows * n_cols):
        row = i // n_cols
        col = i % n_cols
        axes[row, col].axis('off')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    plt.show()


def plot_metrics_comparison(metrics_data: Dict[str, Dict[str, float]],
                           metric_name: str = "Dice",
                           save_path: Optional[Union[str, Path]] = None) -> None:
    """
    绘制指标对比图
    
    Args:
        metrics_data: 指标数据 {method: {metric: value}}
        metric_name: 要对比的指标名称
        save_path: 保存路径
    """
    setup_matplotlib()
    
    methods = list(metrics_data.keys())
    values = [metrics_data[method].get(metric_name, 0) for method in methods]
    
    fig, ax = plt.subplots(1, 1, figsize=(10, 6))
    
    bars = ax.bar(methods, values, color='skyblue', edgecolor='navy', alpha=0.7)
    
    # 添加数值标签
    for bar, value in zip(bars, values):
        height = bar.get_height()
        ax.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                f'{value:.3f}', ha='center', va='bottom', fontweight='bold')
    
    ax.set_ylabel(f'{metric_name} Score')
    ax.set_title(f'不同方法的{metric_name}指标对比')
    ax.set_ylim(0, 1.1)
    ax.grid(True, alpha=0.3)
    
    plt.xticks(rotation=45)
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    plt.show()


def create_attention_heatmap(attention_weights: np.ndarray,
                           image: np.ndarray,
                           title: str = "Attention Heatmap",
                           save_path: Optional[Union[str, Path]] = None) -> None:
    """
    创建注意力热力图
    
    Args:
        attention_weights: 注意力权重 (H, W)
        image: 原始图像 (H, W, C) 或 (H, W)
        title: 图像标题
        save_path: 保存路径
    """
    setup_matplotlib()
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
    
    # 原图
    if len(image.shape) == 3:
        ax1.imshow(image)
    else:
        ax1.imshow(image, cmap='gray')
    ax1.set_title('原始图像')
    ax1.axis('off')
    
    # 注意力热力图
    im = ax2.imshow(attention_weights, cmap='hot', interpolation='nearest')
    ax2.set_title('注意力热力图')
    ax2.axis('off')
    
    # 添加颜色条
    plt.colorbar(im, ax=ax2, fraction=0.046, pad=0.04)
    
    plt.suptitle(title, fontsize=14)
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    plt.show()


def plot_histogram(data: np.ndarray,
                  bins: int = 50,
                  title: str = "Histogram",
                  xlabel: str = "Value",
                  ylabel: str = "Frequency",
                  save_path: Optional[Union[str, Path]] = None) -> None:
    """
    绘制直方图
    
    Args:
        data: 数据
        bins: 直方图箱数
        title: 图表标题
        xlabel: X轴标签
        ylabel: Y轴标签
        save_path: 保存路径
    """
    setup_matplotlib()
    
    fig, ax = plt.subplots(1, 1, figsize=(8, 6))
    
    ax.hist(data, bins=bins, alpha=0.7, color='skyblue', edgecolor='black')
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    plt.show()


def save_prediction_visualization(image: np.ndarray,
                                 ground_truth: np.ndarray,
                                 prediction: np.ndarray,
                                 save_path: Union[str, Path],
                                 title: str = "Prediction Visualization") -> None:
    """
    保存预测结果可视化
    
    Args:
        image: 原始图像
        ground_truth: 真实掩码
        prediction: 预测掩码
        save_path: 保存路径
        title: 图像标题
    """
    setup_matplotlib()
    
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    
    # 原图
    if len(image.shape) == 3:
        axes[0].imshow(image)
    else:
        axes[0].imshow(image, cmap='gray')
    axes[0].set_title('原始图像')
    axes[0].axis('off')
    
    # 真实掩码
    axes[1].imshow(ground_truth, cmap='Reds')
    axes[1].set_title('真实掩码')
    axes[1].axis('off')
    
    # 预测掩码
    axes[2].imshow(prediction, cmap='Blues')
    axes[2].set_title('预测掩码')
    axes[2].axis('off')
    
    plt.suptitle(title, fontsize=14)
    plt.tight_layout()
    
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()

